import jax
import jax.numpy as jnp
import equinox as eqx
from jaxtyping import Array, Float, Int, Bool, PyTree
from .modules.utils import forward_pass, get_dts # , pad_to_len
from .modules.lnm import lnm_logp, lnm_sample

class GRU(eqx.Module):
    cell: eqx.Module
    hdim: int

    def __init__(self, hdim: int, embed_dim: int, key: Array):
        self.cell = eqx.nn.GRUCell(embed_dim+1, hdim, key=key)
        self.hdim = hdim

    def step(self, h: Float[Array, "hdim"], embed: Float[Array, "embed_dim"], dt: Float[Array, ""]):
        log_dt = jnp.log(dt.clip(1e-6))
        x = jnp.concatenate((embed, log_dt[None])) # (embed_dim+1, )
        h = self.cell(x, h)
        return h
    
    def __call__(self, dts: Float[Array, "T"], embeds: Float[Array, "T embed_dim"]):
        h_init = jnp.zeros(self.hdim)
        xs = (dts, embeds)
        def f(h, x): # This is the step function for scan. Takes (carry, input). Gives (carry, output)
            dt, embed = x
            h = self.step(h, embed, dt)
            return h, h
        carry, hs = jax.lax.scan(f, h_init, xs)
        return carry, hs
    

class LogNormMix(eqx.Module):
    enc: GRU
    stack: eqx.Module
    embed: eqx.Module
    num_component: int
    num_types: int
    hdim: int

    def __init__(self, hdim: int, num_types: int, embed_dim: int, num_components: int, key: Array):
        key1, key2, key3 = jax.random.split(key, 3)
        self.enc = GRU(hdim, embed_dim, key1)
        self.stack = eqx.nn.Linear(hdim, 3*num_components+num_types, key=key2)
        self.embed = eqx.nn.Embedding(num_types, embed_dim, key=key3)
        self.num_component = num_components
        self.num_types = num_types
        self.hdim = hdim

    @eqx.filter_jit
    def __call__(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], key=None):
        dts = get_dts(ts)
        dts = jnp.where(mask, dts, 0.)
        embeds = jax.vmap(self.embed)(marks)
        _, hs = self.enc(dts, embeds)
        ll, dt_ll, mark_ll = jax.vmap(self._get_ll)\
                    (hs[:-1], 
                     dts[1:], 
                     marks[1:], 
                     mask[1:]) # (T-1, )
        ll = ll.sum()
        dt_ll = dt_ll.sum()
        mark_ll = mark_ll.sum()
        return ll, dt_ll, mark_ll

    def _get_ll(self, h: Float[Array, "hdim"], dt: Float[Array, ""], mark: Int[Array, ""], mask: Bool[Array, ""]):
        locs, log_scales, log_weights, mark_logits = self._get_dist(h)
        dt_ll = lnm_logp(dt, locs, log_scales, log_weights)
        mark_ll = mark_logits[mark]
        dt_ll = jnp.where(mask, dt_ll, 0.)
        mark_ll = jnp.where(mask, mark_ll, 0.)
        ll = dt_ll + mark_ll
        return ll, dt_ll, mark_ll

    @eqx.filter_jit
    def _get_dist(self, h: Float[Array, 'hdim']):
        param = self.stack(h) # (3*num_components+num_types)
        t_param = param[:3*self.num_component]
        mark_param = param[3*self.num_component:]
        locs, log_scales, log_weights = jnp.split(t_param, 3, -1)
        log_weights = jax.nn.log_softmax(log_weights, axis=-1) # (num_components, )
        mark_logits = jax.nn.log_softmax(mark_param, -1)
        return locs, log_scales, log_weights, mark_logits
    
    def _predict(self, h: Float[Array, "hdim"], dt_max: Float[Array, ""]):
        locs, log_scales, log_weights, mark_logits = self._get_dist(h) # (num_components, )
        means = jnp.exp(locs + jnp.exp(log_scales)**2/2) # (num_components, )
        means_weight = means * jnp.exp(log_weights) # (num_components, )
        Efdt = means_weight.sum() # ()
        Efdt = jnp.clip(Efdt, max=dt_max)
        mark_predict = mark_logits.argmax(-1) # ()
        return Efdt, mark_predict
    
    @eqx.filter_jit
    def rolling_predict(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], dt_max: float, key=None):
        dts = get_dts(ts)
        dts = jnp.where(mask, dts, 0.)
        embeds = jax.vmap(self.embed)(marks)
        _, hs = self.enc(dts, embeds)
        dt_predict, mark_predict = jax.vmap(self._predict, (0, None))(hs[:-1], dt_max)
        return (dt_predict, mark_predict), (dts[1:], marks[1:]), mask[1:]